﻿#include <iostream>
#include <vector>
#include <cmath>
#include <iomanip>
#include "../include/vector.h"
#include "../include/convolution.h"
#include "../include/iir.h"

using namespace std;
using namespace splab;

// 生成测试信号：正弦波加高频噪声
template<typename Type>
Vector<Type> generateTestSignal(int length, Type fs) {
    Vector<Type> signal(length);
    Type t;
    
    for(int i = 0; i < length; i++) {
        t = Type(i) / fs;
        // 生成1Hz的正弦波加上10Hz的高频噪声
        signal[i] = sin(2 * PI * t) + 0.5 * sin(2 * PI * 10 * t);
    }
    
    return signal;
}

// 应用IIR滤波器
Vector<double> applyFilter(const Vector<double>& signal, const Vector<double>& numCoefs, const Vector<double>& denCoefs) {
    int signalLength = signal.size();
    Vector<double> filteredSignal(signalLength);
    
    // 初始化缓存
    vector<double> x_history(numCoefs.size(), 0.0);
    vector<double> y_history(denCoefs.size() - 1, 0.0);
    
    // 应用滤波器
    for(int n = 0; n < signalLength; n++) {
        // 更新输入历史
        for(int i = x_history.size() - 1; i > 0; i--) {
            x_history[i] = x_history[i-1];
        }
        x_history[0] = signal[n];
        
        // 计算输出
        double y = 0.0;
        for(int i = 0; i < numCoefs.size(); i++) {
            y += numCoefs[i] * x_history[i];
        }
        for(int i = 0; i < y_history.size(); i++) {
            y -= denCoefs[i+1] * y_history[i];
        }
        y /= denCoefs[0];
        
        // 更新输出历史
        for(int i = y_history.size() - 1; i > 0; i--) {
            y_history[i] = y_history[i-1];
        }
        if(!y_history.empty()) {
            y_history[0] = y;
        }
        
        filteredSignal[n] = y;
    }
    
    return filteredSignal;
}

int main() {
    // 设置采样频率和信号长度
    double fs = 100.0;  // 采样频率100Hz
    int signalLength = 1000;  // 信号长度1000点
    
    // 生成测试信号
    Vector<double> signal = generateTestSignal<double>(signalLength, fs);
    
    // 设计低通IIR滤波器（Butterworth）
    IIR lowpassFilter("lowpass", "Butterworth");
    lowpassFilter.setParams(fs, 2.0, -3.0, 5.0, -40.0);  // 截止频率2Hz
    lowpassFilter.design();
    
    // 获取低通滤波器系数
    Vector<double> lpNumCoefs = lowpassFilter.getNumCoefs();  // 分子系数
    Vector<double> lpDenCoefs = lowpassFilter.getDenCoefs();  // 分母系数
    
    // 设计高通IIR滤波器（Butterworth）
    IIR highpassFilter("highpass", "Butterworth");
    highpassFilter.setParams(fs, 8.0, -3.0, 5.0, -40.0);  // 截止频率8Hz
    highpassFilter.design();
    
    // 获取高通滤波器系数
    Vector<double> hpNumCoefs = highpassFilter.getNumCoefs();
    Vector<double> hpDenCoefs = highpassFilter.getDenCoefs();
    
    // 应用滤波器
    Vector<double> lowpassFiltered = applyFilter(signal, lpNumCoefs, lpDenCoefs);
    Vector<double> highpassFiltered = applyFilter(signal, hpNumCoefs, hpDenCoefs);
    
    // 输出结果
    cout << "原始信号、低通滤波和高通滤波后信号的前20个采样点：" << endl;
    cout << setw(10) << "时间(s)" 
         << setw(15) << "原始信号" 
         << setw(15) << "低通滤波" 
         << setw(15) << "高通滤波" << endl;
    cout << "--------------------------------------------------------" << endl;
    
    for(int i = 0; i < 20; i++) {
        cout << setw(10) << fixed << setprecision(3) << i/fs
             << setw(15) << signal[i]
             << setw(15) << lowpassFiltered[i]
             << setw(15) << highpassFiltered[i] << endl;
    }
    
    return 0;
}
